
import numpy as np
import math
import random
from datetime import datetime

def generate_square_odometry(side_length=1000, velocity=1.0, sample_rate=10):
    """生成1km正方形闭环轨迹的里程计数据"""
    timestamps = []
    positions = []
    angles = []
    
    # 轨迹参数计算 (单位：米)
    segment_length = side_length / 4
    segment_time = segment_length / velocity
    total_points = int(segment_time * sample_rate * 4)
    
    # 四段轨迹：北→东→南→西
    for t in np.linspace(0, segment_time*4, total_points):
        timestamps.append(t)
        
        # 分段处理
        if t < segment_time:  # 北向移动
            x = -side_length/2
            y = -side_length/2 + velocity*t
            theta = math.pi/2
        elif t < segment_time*2:  # 东向移动
            x = -side_length/2 + velocity*(t-segment_time)
            y = side_length/2
            theta = 0
        elif t < segment_time*3:  # 南向移动
            x = side_length/2
            y = side_length/2 - velocity*(t-segment_time*2)
            theta = -math.pi/2
        else:  # 西向移动
            x = side_length/2 - velocity*(t-segment_time*3)
            y = -side_length/2
            theta = math.pi
        
        # 添加里程计噪声 (5%误差)
        x += random.gauss(0, 0.05*velocity)
        y += random.gauss(0, 0.05*velocity)
        theta += random.gauss(0, 0.02)
        
        positions.append((x, y))
        angles.append(theta)
    
    return timestamps, positions, angles

def generate_lidar_scan(position, angle, map_size=1050):
    """生成带噪声的激光雷达扫描数据"""
    scan = []
    num_beams = 360  # 360度全覆盖
    max_range = 50.0  # 50米最大测距
    
    for beam_angle in np.linspace(0, 2*math.pi, num_beams, endpoint=False):
        absolute_angle = angle + beam_angle
        hit_distance = max_range
        
        # 检测正方形边界
        x0, y0 = position
        for d in np.linspace(0, max_range, 50):
            x = x0 + d * math.cos(absolute_angle)
            y = y0 + d * math.sin(absolute_angle)
            
            if abs(x) > map_size/2 or abs(y) > map_size/2:
                hit_distance = d
                break
        
        # 添加测量噪声 (1%误差 + 随机异常值)
        if random.random() > 0.95:
            hit_distance = random.uniform(0, max_range)
        else:
            hit_distance *= (1 + random.gauss(0, 0.01))
        
        scan.append(min(max_range, hit_distance))
    
    return scan

def save_cartographer_data(filename, timestamps, odom_data, lidar_data):
    """保存为Cartographer兼容格式"""
    with open(filename, 'w') as f:
        f.write("# Format: timestamp,x,y,theta,range1,range2,...\n")
        for t, (x, y, theta), scan in zip(timestamps, odom_data, lidar_data):
            #f.write(f"{t:.6f},{x:.6f},{y:.6f},{theta:.6f},")
            f.write("{0:.6f},{1:.6f},{2:.6f},{3:.6f},".format(t, x, y, theta))
            f.write(",".join(["{:.6f}".format(r) for r in scan]) + "\n")

if __name__ == "__main__":
    # 生成1km正方形轨迹数据
    timestamps, positions, angles = generate_square_odometry()
    odom_data = list(zip([p[0] for p in positions], 
                        [p[1] for p in positions], 
                        angles))
    lidar_data = [generate_lidar_scan(p, a) for p, a in zip(positions, angles)]
    
    # 保存数据集
    #filename = f"square_1km_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    filename = "square_1km_{}.csv".format(datetime.now().strftime('%Y%m%d_%H%M%S'))
    save_cartographer_data(filename, timestamps, odom_data, lidar_data)
    #print(f"Generated 1km square dataset with {len(timestamps)} frames")
    #print(f"Start/End position: {positions[0]} -> {positions[-1]}")

